import open3d as o3d
import os
from initialize_config import initialize_config
import json
import time

def converter(dir, file, depth_only):
    print(dir, file)
    reader = o3d.io.AzureKinectMKVReader()
    reader.open(os.path.join(dir, file))
    os.makedirs( os.path.join(dir,"out" , "vis","depth"), exist_ok=True)
    if not depth_only:
        meta = reader.get_metadata()
        o3d.io.write_azure_kinect_mkv_metadata(os.path.join(dir,"out" , "vis", "intrinsic.json"), meta)
        config = {
            "path_dataset": os.path.join(dir,"out" , "vis"),
            "path_intrinsic": os.path.join(dir,"out" , "vis", "intrinsic.json")
        }
        initialize_config(config)
        with open(os.path.join(dir,"out" , "vis", "config.json"), "w") as f:
            json.dump(config, f, indent=4)
        os.makedirs( os.path.join(dir,"out" , "vis","color"), exist_ok=True)
        
    idx = 0
    while not reader.is_eof():
        rgbd = reader.next_frame()
        if rgbd is None:
            continue
        depth_filename = os.path.join(dir,"out" , "vis","depth", '{:05d}.png'.format(idx))
        if not depth_only:
            color_filename = os.path.join(dir,"out" , "vis","color", '{:05d}.jpg'.format(idx))
            o3d.io.write_image(color_filename, rgbd.color)
        o3d.io.write_image(depth_filename, rgbd.depth)
        idx += 1
    reader.close()
    with open(os.path.join(dir, "updated"), "w"):
        pass
    
if __name__ == "__main__":
    DIR = "D:\\mkv\\azure5"
    while(True):
        try:
            for day in os.listdir(DIR):
                for clip in os.listdir(os.path.join(DIR, day)):
                    if not os.path.isdir(os.path.join(DIR, day, clip)):
                        continue
                    files = os.listdir(os.path.join(DIR, day, clip))
                    for file in files:
                        if "updated" in files:
                            continue
                        if file.endswith(".mkv"):
                            if "vis" in files:
                                converter(os.path.join(DIR, day, clip), file, True)
                            elif "depth" in files:
                                converter(os.path.join(DIR, day, clip), file, False)
                        break
        except Exception as e:
            print(str(e))
        finally:
            time.sleep(5)
            